import csv
import multiprocessing

import torch
from eval import show_without_plt
from eval import check_orthogonality_and_determinant
from optimizer import SGDG
from ffn_getter.get_each_proj import save_proj

GD_W = 1e-2
cayley_alpha = 1e-2
loss_factor = 1e-2

csv_name = "12-30"

# 选择计算设备：GPU或CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化随机正交矩阵（旋转矩阵）
def random_orthogonal_matrix(dim, device):
    """生成一个随机的正交矩阵，且行列式值为 1"""
    matrix = torch.randn(dim, dim, device=device)  # 随机生成一个矩阵
    Q, _ = torch.linalg.qr(matrix)  # 使用QR分解生成正交矩阵
    # 修正行列式为 1
    if torch.det(Q) < 0:
        Q[:, 0] = -Q[:, 0]
    return Q


class rotateMatrix(torch.nn.Module):
    def __init__(self, input_dim, device):
        super(rotateMatrix, self).__init__()
        self.input_dim = input_dim
        # self.R = random_orthogonal_matrix(input_dim, device).requires_grad_()  # 确保R可以计算梯度

        self.R = torch.nn.Parameter(random_orthogonal_matrix(input_dim, device))

    def forward(self, X):
        # 旋转输入矩阵 X
        rotated_X = torch.matmul(X, self.R)
        return rotated_X


def compute_loss_l2(R, X, l2_factor=loss_factor):
    """计算 L2 范数损失（基于 R * X）"""
    rotated_X = torch.matmul(X, R)  # 与 R 进行矩阵乘法

    l2_loss = torch.sum(rotated_X ** 2)  # L2 范数损失（平方和）
    total_loss = l2_factor * l2_loss  # 总损失为 L2 范数损失

    return total_loss
def compute_loss_l1(R, X, l1_factor=loss_factor):
    """计算 L1 范数损失（基于 R * X）"""
    rotated_X = torch.matmul(X, R)  # 与R进行矩阵乘法

    l1_loss = torch.sum(torch.abs(rotated_X))  # L1正则化项
    total_loss = l1_factor * l1_loss  # 总损失为 L1 范数损失
    return total_loss

dataset_start = 0
range_limit = 32

def load_proj():
    global dataset_start
    global range_limit
    save_proj(dataset_num, dataset_num + 1)
    dataset_num += 1
    return torch.stack([torch.load(f'ffn/second_line_input/tensor{i}.pth', weights_only=True).to(torch.float32) for i in range(range_limit)]).to(device)

# process = multiprocessing.Process(target=run_save_proj_in_process)
# process.start()
# process.join()

load_fms = load_proj()

rm_dim = load_fms[0].shape[2]
print(rm_dim)

model = rotateMatrix(rm_dim, device).to(device)

optimizer = SGDG(model.parameters(), lr=GD_W,stiefel=True)

fieldnames = ['epoch'] + list( show_without_plt(load_fms[0], print_info=False).keys() ) + ["orthogonal_error", "det"]

def csv_creater(file):
    global fieldnames
    writer = csv.DictWriter(file, fieldnames=fieldnames)
    writer.writeheader()
    file.flush()
    return writer

saved = False

def save_log(epoch, input, orthogonal_error, file,writer):
    global fieldnames
    global saved
    print("------------------------------------")
    print(f"epoch: {epoch}", end='')
    result = show_without_plt(input, print_info=True)
    print(f"orthogonal_error: {orthogonal_error[0]}")
    print(f"det: {orthogonal_error[1]}")
    result["epoch"] = epoch
    result["orthogonal_error"] = orthogonal_error[0]
    result["det"] = orthogonal_error[1]
    writer.writerow(result)
    # save_csv = file.write(','.join([str(i) for i in arr]) + '\n')
    torch.save(model.R, f'{csv_name}.pth')

    if result["lower_ratio_1e_4"] > 0.2 and saved == False:
        torch.save(model.R, f'{csv_name}_target_{epoch}.pth')
        print("save target, ratio > 0.2, epoch: ",epoch)
        saved = True

    file.flush()
    print("------------------------------------")



print("start training")

file = open(f'{csv_name}.csv', mode='w', newline='')

writer_out = csv_creater(file)

with torch.no_grad():
    save_log(-1, load_fms, check_orthogonality_and_determinant(model.R), file,writer_out )
for epoch in range(2000):

    print(f"epoch: {epoch}")
    for i in range(range_limit):
        print(f"i: {i}",end=' ')
        input_data = load_fms[i]
        optimizer.zero_grad()

        output = model(input_data)

        total_loss= compute_loss_l1(model.R, input_data)

        total_loss.backward()

        optimizer.step()
    with torch.no_grad():
        orthogonal_error = check_orthogonality_and_determinant(model.R)

        # 使用 show_without_plt 显示旋转后的矩阵（非plt版本）
    if epoch % 3 == 0:
        with torch.no_grad():
            save_log(epoch, load_fms @ model.R, orthogonal_error, file,writer_out)

    # if epoch % 200 == 0 and epoch != 0:
    #     load_fms = load_proj()

file.close()



